import oneflow as torch
import oneflow.nn as nn

def get_length_mask(tensor, tensor_length):
    b, t, _ = tensor.size()  
    mask = tensor.new_zeros([b, t], dtype=torch.uint8)
    for i, length in enumerate(tensor_length):
        length = length.item()
        mask[i].narrow(0, 0, length).fill_(1)
    return mask > 0


def get_local_mask(tensor, length, left_context, right_context):
    b, t, _ = tensor.size()

    if left_context + 1 > t:
        left_context = t - 1

    if left_context >= 0:
        left_mask = torch.tril(torch.ones((t, t), dtype=torch.uint8), diagonal=-(left_context + 1))
    else:
        left_mask = torch.zeros((t, t), dtype=torch.uint8)

    if right_context >= 0:
        right_mask = torch.triu(torch.ones((t, t), dtype=torch.uint8), diagonal=right_context + 1)
    else:
        right_mask = torch.zeros((t, t), dtype=torch.uint8)

    mask = left_mask + right_mask
    mask = mask.unsqueeze(0).repeat(b, 1, 1) < 1
    mask = mask.to(tensor.device)

    for b in range(tensor.size(0)):
        mask[b, length[b].item():] = True
        mask[b, :, length[b].item():] = False

    return mask.to(tensor.device)


def init_weights(m, mode='xavier_uniform'):


    if isinstance(m, (nn.Conv1d, nn.Linear)):
        if mode == 'xavier_uniform':
            nn.init.xavier_uniform_(m.weight, gain=1.0)
        elif mode == 'xavier_normal':
            nn.init.xavier_normal_(m.weight, gain=1.0)
        elif mode == 'kaiming_uniform':
            nn.init.kaiming_uniform_(m.weight, nonlinearity="relu")
        elif mode == 'kaiming_normal':
            nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
        else:
            raise ValueError("Unknown Initialization mode: {0}".format(mode))
    elif isinstance(m, nn.BatchNorm1d):
        if m.track_running_stats:
            m.running_mean.zero_()
            m.running_var.fill_(1)
            m.num_batches_tracked.zero_()
        if m.affine:
            nn.init.ones_(m.weight)
            nn.init.zeros_(m.bias)